"""
04_conditional_routing.py - Conditional Routing and Control Flow
================================================================
Chapter 6: Making your graph smart with dynamic decisions.

KEY CONCEPTS:
  ★ Conditional Edges - Choose next node based on state
  ★ Router Function - Returns the NAME of the next node
  ★ Literal type hint - Constrains router return values
  ★ Cycles/Loops - The agent pattern (decide -> act -> repeat)

Run: python src/04_conditional_routing.py
"""

from typing import TypedDict, Literal
from langgraph.graph import StateGraph, START, END


class RouterState(TypedDict):
    query: str
    query_type: str
    response: str
    attempts: int


def classify_query(state: RouterState) -> dict:
    """Classify the incoming query type."""
    query = state["query"].lower()
    if "code" in query or "python" in query or "function" in query:
        return {"query_type": "coding"}
    elif "math" in query or "calculate" in query or "compute" in query:
        return {"query_type": "math"}
    elif "write" in query or "essay" in query or "draft" in query:
        return {"query_type": "writing"}
    else:
        return {"query_type": "general"}


def handle_coding(state: RouterState) -> dict:
    return {"response": f"🖥️ [Code Agent] Here's a code solution for: {state['query']}"}

def handle_math(state: RouterState) -> dict:
    return {"response": f"🔢 [Math Agent] Computing result for: {state['query']}"}

def handle_writing(state: RouterState) -> dict:
    return {"response": f"✍️ [Writing Agent] Drafting content for: {state['query']}"}

def handle_general(state: RouterState) -> dict:
    return {"response": f"💬 [General Agent] Responding to: {state['query']}"}


# ============================================================
# THE ROUTER FUNCTION
# ============================================================
# Returns the NAME of the next node to execute.
# The Literal type hint lists all valid return values.

def route_query(state: RouterState) -> Literal[
    "coding_handler", "math_handler", "writing_handler", "general_handler"
]:
    """Route to the appropriate handler based on query type.
    
    ★ IMPORTANT: Return value MUST match a node name!
    """
    routing_map = {
        "coding": "coding_handler",
        "math": "math_handler",
        "writing": "writing_handler",
        "general": "general_handler",
    }
    return routing_map.get(state["query_type"], "general_handler")


# Build graph with conditional routing
graph = StateGraph(RouterState)

graph.add_node("classify", classify_query)
graph.add_node("coding_handler", handle_coding)
graph.add_node("math_handler", handle_math)
graph.add_node("writing_handler", handle_writing)
graph.add_node("general_handler", handle_general)

graph.add_edge(START, "classify")

# ★ Conditional edge: after "classify", call route_query
#   to determine which handler to execute next
graph.add_conditional_edges(
    "classify",           # Source node
    route_query,          # Router function
    {                     # Mapping: return value -> node name
        "coding_handler": "coding_handler",
        "math_handler": "math_handler",
        "writing_handler": "writing_handler",
        "general_handler": "general_handler",
    }
)

# All handlers lead to END
for handler in ["coding_handler", "math_handler", "writing_handler", "general_handler"]:
    graph.add_edge(handler, END)

app = graph.compile()


if __name__ == "__main__":
    print("=" * 60)
    print("04: Conditional Routing")
    print("=" * 60)
    
    queries = [
        "Write a Python function for sorting",
        "Calculate the area of a circle with radius 5",
        "Write an essay about climate change",
        "What is the meaning of life?",
    ]
    
    for q in queries:
        result = app.invoke({"query": q})
        print(f"\n  Query: {q}")
        print(f"  Type:  {result['query_type']}")
        print(f"  Route: {result['response']}")
    
    print(f"\n{'=' * 60}")
    print("✅ Conditional routing complete!")
    print("=" * 60)
